"""调试工具模块
提供矩阵计算的调试和性能分析功能
"""
import numpy as np
import time
import traceback
import functools
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import warnings
import sys
import psutil
import gc
from contextlib import contextmanager

class MatrixDebugger:
    """矩阵计算调试器"""
    
    def __init__(self):
        self.debug_info = []
        self.performance_stats = {}
        self.memory_usage = []
        
    def log_operation(self, operation_name: str, 
                     inputs: Dict[str, Any], 
                     outputs: Dict[str, Any],
                     execution_time: float,
                     memory_usage: float) -> None:
        """记录操作信息"""
        log_entry = {
            'timestamp': time.time(),
            'operation': operation_name,
            'inputs': self._serialize_inputs(inputs),
            'outputs': self._serialize_outputs(outputs),
            'execution_time': execution_time,
            'memory_usage': memory_usage
        }
        self.debug_info.append(log_entry)
    
    def _serialize_inputs(self, inputs: Dict[str, Any]) -> Dict[str, str]:
        """序列化输入参数"""
        serialized = {}
        for key, value in inputs.items():
            if isinstance(value, np.ndarray):
                serialized[key] = f"ndarray(shape={value.shape}, dtype={value.dtype})"
            else:
                serialized[key] = str(value)
        return serialized
    
    def _serialize_outputs(self, outputs: Dict[str, Any]) -> Dict[str, str]:
        """序列化输出结果"""
        serialized = {}
        for key, value in outputs.items():
            if isinstance(value, np.ndarray):
                serialized[key] = f"ndarray(shape={value.shape}, dtype={value.dtype})"
            elif isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], np.ndarray):
                shapes = [arr.shape for arr in value]
                serialized[key] = f"list of ndarrays(shapes={shapes})"
            else:
                serialized[key] = str(value)
        return serialized
    
    def get_debug_summary(self) -> Dict[str, Any]:
        """获取调试摘要"""
        if not self.debug_info:
            return {
                "message": "No debug information available",
                "operations_count": 0,
                "errors_count": 0,
                "total_time": 0.0,
                "memory_usage": 0.0
            }
        
        total_operations = len(self.debug_info)
        
        # 安全地计算总时间，处理缺失的execution_time字段
        total_time = 0.0
        error_count = 0
        memory_peak = 0.0
        
        for entry in self.debug_info:
            # 安全地获取execution_time
            if 'execution_time' in entry:
                total_time += entry['execution_time']
            elif 'timestamp' in entry and hasattr(self, '_last_timestamp'):
                # 如果没有execution_time，尝试从时间戳计算
                total_time += max(0, entry['timestamp'] - self._last_timestamp)
            
            # 计算错误数量
            if entry.get('type') == 'error' or 'ERROR' in entry.get('operation', ''):
                error_count += 1
            
            # 安全地获取内存使用
            if 'memory_usage' in entry:
                memory_peak = max(memory_peak, entry['memory_usage'])
        
        avg_time = total_time / total_operations if total_operations > 0 else 0.0
        
        operations_count = {}
        for entry in self.debug_info:
            op = entry.get('operation', 'unknown')
            operations_count[op] = operations_count.get(op, 0) + 1
        
        return {
            'operations_count': total_operations,
            'errors_count': error_count,
            'total_time': total_time,
            'memory_usage': memory_peak,
            'total_operations': total_operations,
            'total_execution_time': total_time,
            'average_execution_time': avg_time,
            'operations_breakdown': operations_count,
            'memory_peak': memory_peak
        }
    
    def clear_debug_info(self) -> None:
        """清除调试信息"""
        self.debug_info.clear()
        self.performance_stats.clear()
        self.memory_usage.clear()
    
    def log_error(self, error_message: str) -> None:
        """记录错误信息"""
        error_entry = {
            'timestamp': time.time(),
            'operation': 'ERROR',
            'type': 'error',
            'message': error_message,
            'traceback': traceback.format_exc(),
            'execution_time': 0.0,  # 错误记录的执行时间设为0
            'memory_usage': 0.0     # 错误记录的内存使用设为0
        }
        self.debug_info.append(error_entry)

# 全局调试器实例
_global_debugger = MatrixDebugger()

def debug_matrix_operation(operation_name: str = None):
    """矩阵操作调试装饰器"""
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            op_name = operation_name or func.__name__
            
            # 记录开始时间和内存
            start_time = time.time()
            start_memory = psutil.Process().memory_info().rss / 1024 / 1024  # MB
            
            try:
                # 执行函数
                result = func(*args, **kwargs)
                
                # 记录结束时间和内存
                end_time = time.time()
                end_memory = psutil.Process().memory_info().rss / 1024 / 1024  # MB
                
                execution_time = end_time - start_time
                memory_delta = end_memory - start_memory
                
                # 准备输入输出信息
                inputs = {f'arg_{i}': arg for i, arg in enumerate(args)}
                inputs.update(kwargs)
                
                if isinstance(result, tuple):
                    outputs = {f'output_{i}': res for i, res in enumerate(result)}
                else:
                    outputs = {'output': result}
                
                # 记录调试信息
                _global_debugger.log_operation(
                    op_name, inputs, outputs, execution_time, memory_delta
                )
                
                return result
                
            except Exception as e:
                # 记录异常信息
                error_info = {
                    'error_type': type(e).__name__,
                    'error_message': str(e),
                    'traceback': traceback.format_exc()
                }
                
                end_time = time.time()
                execution_time = end_time - start_time
                
                inputs = {f'arg_{i}': arg for i, arg in enumerate(args)}
                inputs.update(kwargs)
                
                _global_debugger.log_operation(
                    f"{op_name}_ERROR", inputs, error_info, execution_time, 0
                )
                
                raise
        
        return wrapper
    return decorator

@contextmanager
def matrix_performance_monitor(operation_name: str):
    """性能监控上下文管理器"""
    start_time = time.time()
    start_memory = psutil.Process().memory_info().rss / 1024 / 1024
    
    try:
        yield
    finally:
        end_time = time.time()
        end_memory = psutil.Process().memory_info().rss / 1024 / 1024
        
        execution_time = end_time - start_time
        memory_delta = end_memory - start_memory
        
        print(f"Operation '{operation_name}':")
        print(f"  Execution time: {execution_time:.4f} seconds")
        print(f"  Memory delta: {memory_delta:.2f} MB")

class NumericalStabilityChecker:
    """数值稳定性检查器"""
    
    @staticmethod
    def check_condition_number(matrix: np.ndarray, 
                             threshold: float = 1e12) -> Dict[str, Union[float, bool, str]]:
        """检查矩阵条件数"""
        try:
            cond_num = np.linalg.cond(matrix)
            is_well_conditioned = cond_num < threshold
            
            if cond_num == np.inf:
                status = "singular"
            elif cond_num > threshold:
                status = "ill-conditioned"
            else:
                status = "well-conditioned"
            
            return {
                'condition_number': float(cond_num),
                'is_well_conditioned': is_well_conditioned,
                'status': status,
                'threshold': threshold
            }
        except Exception as e:
            return {
                'error': str(e),
                'status': 'error'
            }
    
    @staticmethod
    def check_numerical_rank(matrix: np.ndarray, 
                           tolerance: float = None) -> Dict[str, Union[int, float, bool]]:
        """检查数值秩"""
        if tolerance is None:
            tolerance = max(matrix.shape) * np.finfo(matrix.dtype).eps * np.linalg.norm(matrix)
        
        try:
            singular_values = np.linalg.svd(matrix, compute_uv=False)
            numerical_rank = np.sum(singular_values > tolerance)
            theoretical_rank = min(matrix.shape)
            is_full_rank = numerical_rank == theoretical_rank
            
            return {
                'numerical_rank': int(numerical_rank),
                'theoretical_max_rank': int(theoretical_rank),
                'is_full_rank': is_full_rank,
                'tolerance': float(tolerance),
                'smallest_singular_value': float(np.min(singular_values)),
                'largest_singular_value': float(np.max(singular_values))
            }
        except Exception as e:
            return {
                'error': str(e)
            }
    
    @staticmethod
    def check_eigenvalue_stability(matrix: np.ndarray) -> Dict[str, Union[bool, float, List[float]]]:
        """检查特征值稳定性"""
        if matrix.shape[0] != matrix.shape[1]:
            return {'error': 'Matrix must be square for eigenvalue analysis'}
        
        try:
            eigenvalues = np.linalg.eigvals(matrix)
            real_parts = np.real(eigenvalues)
            imag_parts = np.imag(eigenvalues)
            
            # 检查稳定性（所有特征值实部 < 0）
            is_stable = np.all(real_parts < 0)
            
            # 检查是否在单位圆内（离散系统）
            magnitudes = np.abs(eigenvalues)
            is_discrete_stable = np.all(magnitudes < 1)
            
            return {
                'is_continuous_stable': is_stable,
                'is_discrete_stable': is_discrete_stable,
                'max_real_part': float(np.max(real_parts)),
                'max_magnitude': float(np.max(magnitudes)),
                'eigenvalues_real': real_parts.tolist(),
                'eigenvalues_imag': imag_parts.tolist()
            }
        except Exception as e:
            return {
                'error': str(e)
            }

class MatrixValidator:
    """矩阵验证器"""
    
    @staticmethod
    def validate_matrix_properties(matrix: np.ndarray) -> Dict[str, bool]:
        """验证矩阵基本属性"""
        properties = {}
        
        # 基本检查
        properties['is_finite'] = np.all(np.isfinite(matrix))
        properties['has_nan'] = np.any(np.isnan(matrix))
        properties['has_inf'] = np.any(np.isinf(matrix))
        
        if matrix.shape[0] == matrix.shape[1]:  # 方阵
            properties['is_square'] = True
            properties['is_symmetric'] = np.allclose(matrix, matrix.T)
            properties['is_hermitian'] = np.allclose(matrix, matrix.conj().T)
            properties['is_positive_definite'] = MatrixValidator._is_positive_definite(matrix)
            properties['is_positive_semidefinite'] = MatrixValidator._is_positive_semidefinite(matrix)
            properties['is_orthogonal'] = MatrixValidator._is_orthogonal(matrix)
            properties['is_unitary'] = MatrixValidator._is_unitary(matrix)
        else:
            properties['is_square'] = False
        
        return properties
    
    @staticmethod
    def _is_positive_definite(matrix: np.ndarray) -> bool:
        """检查是否正定"""
        try:
            np.linalg.cholesky(matrix)
            return True
        except np.linalg.LinAlgError:
            return False
    
    @staticmethod
    def _is_positive_semidefinite(matrix: np.ndarray) -> bool:
        """检查是否半正定"""
        try:
            eigenvalues = np.linalg.eigvals(matrix)
            return np.all(eigenvalues >= -1e-10)  # 允许小的数值误差
        except:
            return False
    
    @staticmethod
    def _is_orthogonal(matrix: np.ndarray) -> bool:
        """检查是否正交"""
        try:
            product = matrix @ matrix.T
            identity = np.eye(matrix.shape[0])
            return np.allclose(product, identity)
        except:
            return False
    
    @staticmethod
    def _is_unitary(matrix: np.ndarray) -> bool:
        """检查是否酉矩阵"""
        try:
            product = matrix @ matrix.conj().T
            identity = np.eye(matrix.shape[0])
            return np.allclose(product, identity)
        except:
            return False

def get_debug_summary() -> Dict[str, Any]:
    """获取全局调试摘要"""
    return _global_debugger.get_debug_summary()

def clear_debug_info() -> None:
    """清除全局调试信息"""
    _global_debugger.clear_debug_info()

def print_matrix_info(matrix: np.ndarray, name: str = "Matrix") -> None:
    """打印矩阵详细信息"""
    print(f"\n=== {name} Information ===")
    print(f"Shape: {matrix.shape}")
    print(f"Data type: {matrix.dtype}")
    print(f"Size: {matrix.size} elements")
    print(f"Memory usage: {matrix.nbytes / 1024 / 1024:.2f} MB")
    
    # 统计信息
    print(f"\nStatistics:")
    print(f"  Min: {np.min(matrix):.6f}")
    print(f"  Max: {np.max(matrix):.6f}")
    print(f"  Mean: {np.mean(matrix):.6f}")
    print(f"  Std: {np.std(matrix):.6f}")
    
    # 特殊值检查
    nan_count = np.sum(np.isnan(matrix))
    inf_count = np.sum(np.isinf(matrix))
    zero_count = np.sum(matrix == 0)
    
    print(f"\nSpecial values:")
    print(f"  NaN count: {nan_count}")
    print(f"  Inf count: {inf_count}")
    print(f"  Zero count: {zero_count}")
    print(f"  Sparsity: {zero_count / matrix.size:.2%}")
    
    # 方阵特性
    if matrix.shape[0] == matrix.shape[1]:
        print(f"\nSquare matrix properties:")
        try:
            det = np.linalg.det(matrix)
            trace = np.trace(matrix)
            cond = np.linalg.cond(matrix)
            rank = np.linalg.matrix_rank(matrix)
            
            print(f"  Determinant: {det:.6f}")
            print(f"  Trace: {trace:.6f}")
            print(f"  Condition number: {cond:.2e}")
            print(f"  Rank: {rank}")
            
            # 矩阵性质
            properties = MatrixValidator.validate_matrix_properties(matrix)
            print(f"  Symmetric: {properties.get('is_symmetric', False)}")
            print(f"  Positive definite: {properties.get('is_positive_definite', False)}")
            print(f"  Orthogonal: {properties.get('is_orthogonal', False)}")
            
        except Exception as e:
            print(f"  Error computing properties: {e}")
    
    print("=" * 40)

def benchmark_operation(func: Callable, *args, iterations: int = 10, **kwargs) -> Dict[str, float]:
    """基准测试函数"""
    times = []
    memory_usage = []
    
    for _ in range(iterations):
        # 垃圾回收
        gc.collect()
        
        start_memory = psutil.Process().memory_info().rss / 1024 / 1024
        start_time = time.time()
        
        try:
            result = func(*args, **kwargs)
        except Exception as e:
            return {'error': str(e)}
        
        end_time = time.time()
        end_memory = psutil.Process().memory_info().rss / 1024 / 1024
        
        times.append(end_time - start_time)
        memory_usage.append(end_memory - start_memory)
    
    return {
        'mean_time': np.mean(times),
        'std_time': np.std(times),
        'min_time': np.min(times),
        'max_time': np.max(times),
        'mean_memory_delta': np.mean(memory_usage),
        'iterations': iterations
    }